PolicyEngine-UK compatibility: Loosen SS tolerances and add Numba optimizations#1082
PolicyEngine-UK compatibility: Loosen SS tolerances and add Numba optimizations#1082vahid-ahmadi wants to merge 3 commits intoPSLmodels:masterfrom
Conversation
- Loosen mindist_SS from 1e-9 to 1e-3 for faster GE loop convergence - Loosen RC_SS from 1e-8 to 1e-4 for resource constraint check - Add optional Numba JIT compilation for hot paths: - household.py: marg_ut_cons, marg_ut_labor - tax.py: ETR_wealth, MTR_wealth - txfunc.py: DEP tax rate calculations - Numba is optional - falls back to NumPy if not installed 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
vahid-ahmadi
changed the title
|
These are parameter changes right? Why do we need to change the defaults, we can just change them in og-uk no? |
I started working on Numba JIT optimizations for hot-path functions (compiling Python loops to fast machine code). Still, after benchmarking, most of the speedup is actually coming from the changes to the tolerance parameters (mindist_SS, RC_SS). I can move those parameter changes to OG-UK via p.update_specifications() instead of changing OG-Core defaults. |
|
@vahid-ahmadi Thanks for thinking about improvements to the solution of OG-Core. Can you report some diffs of the Numba changes only (i.e., keep the default tolerances and run with only the numba changes). What is the speed up? These are pretty small loops (and sometimes vectorized operations) so I'm skeptical there would be much improvement. |
Keep mindist_SS (1e-9) and RC_SS (1e-8) at original defaults. Looser tolerances for faster runs should be set via p.update_specifications() in downstream packages like OG-UK. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Merge master's tax_filer changes with Numba optimizations: - tax.py: MTR_wealth now uses tax_filer in both Numba and non-Numba paths - txfunc.py: Keep Numba JIT functions alongside master's updates Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
@vahid-ahmadi Thanks for addressing the concerns we raised. Can you let us know what the benchmark results look like with the current changes? |
3 participants
Performance Optimizations for Steady State Solver
Summary
Changes
1. Solver Tolerance Adjustments (
default_parameters.json)mindist_SS1e-91e-3RC_SS1e-81e-4Benchmark results:
2. Optional Numba JIT Compilation
Added Numba-optimized versions of hot-path functions:
household.py:marg_ut_cons,marg_ut_labortax.py:ETR_wealth,MTR_wealthtxfunc.py: DEP tax rate calculationsFeatures:
@njit(cache=True)for compilation cachingprangewhere applicableTest Plan
🤖 Generated with Claude Code